import os
import shutil
from pathlib import Path
import glob

def clear_temp_directory():
    """清空temp目录"""
    temp_dir = "temp"
    if os.path.exists(temp_dir):
        print(f"清空temp目录: {temp_dir}")
        # 删除temp目录中的所有文件
        for file in os.listdir(temp_dir):
            file_path = os.path.join(temp_dir, file)
            try:
                if os.path.isfile(file_path):
                    os.unlink(file_path)
                    print(f"  已删除: {file}")
            except Exception as e:
                print(f"  删除文件时出错 {file_path}: {e}")
    else:
        print(f"创建temp目录: {temp_dir}")
        os.makedirs(temp_dir)

def copy_samples_to_temp():
    """从样本目录复制图片到temp目录"""
    # 目标目录
    temp_dir = "temp"
    
    # 源目录
    flower_samples_dir = "data/flower/samples"
    face_samples_dir = "data/mini_face/samples"
    
    # 检查源目录是否存在
    if not os.path.exists(flower_samples_dir):
        print(f"错误: 找不到花卉样本目录 {flower_samples_dir}")
        return False
    
    if not os.path.exists(face_samples_dir):
        print(f"错误: 找不到人脸样本目录 {face_samples_dir}")
        return False
    
    # 复制花卉样本
    flower_files = glob.glob(os.path.join(flower_samples_dir, "*.jpg"))
    if not flower_files:
        print(f"警告: 在{flower_samples_dir}中没有找到JPG图片")
    
    for file in flower_files:
        file_name = os.path.basename(file)
        dest_name = f"flower_{file_name}"
        dest_path = os.path.join(temp_dir, dest_name)
        try:
            shutil.copy2(file, dest_path)
            print(f"已复制花卉样本: {file} -> {dest_path}")
        except Exception as e:
            print(f"复制文件时出错 {file}: {e}")
    
    # 复制人脸样本
    face_files = glob.glob(os.path.join(face_samples_dir, "*.jpg"))
    if not face_files:
        print(f"警告: 在{face_samples_dir}中没有找到JPG图片")
    
    for file in face_files:
        file_name = os.path.basename(file)
        dest_name = f"face_{file_name}"
        dest_path = os.path.join(temp_dir, dest_name)
        try:
            shutil.copy2(file, dest_path)
            print(f"已复制人脸样本: {file} -> {dest_path}")
        except Exception as e:
            print(f"复制文件时出错 {file}: {e}")
    
    # 统计结果
    flower_count = len(glob.glob(os.path.join(temp_dir, "flower_*.jpg")))
    face_count = len(glob.glob(os.path.join(temp_dir, "face_*.jpg")))
    
    print(f"\n完成! 共复制了 {flower_count} 张花卉图片和 {face_count} 张人脸图片到 {temp_dir} 目录")
    return True

if __name__ == "__main__":
    print("开始准备样本...")
    clear_temp_directory()
    success = copy_samples_to_temp()
    if success:
        print("样本准备完成,可以运行比较脚本了!")
    else:
        print("样本准备失败!") 